#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
自动化数据集生成脚本
执行指令：
conda activate seal_env && \
python continue_gen_dataset.py \
    --input_path /path/to/source/files \
    --output_path /path/to/output/dataset \
    --train_ratio 0.8

该脚本实现以下功能：
1. 扫描指定路径下的所有文件
2. 图片文件保持原名拷贝
3. Word/PDF文件转换为图片（按文件名_页码.jpg格式命名）
4. 智能选择标注方式：
   - 如果检测到handwriting/printer目录结构，使用特定路径标注逻辑
   - 否则使用PaddleX进行版面检测和自动标注
5. 按比例生成训练和验证数据
6. 输出COCO格式的JSON标注文件

功能说明:
    - 支持常见图片格式：PNG、JPEG、BMP、TIFF等
    - 支持Word文档（.docx）和PDF文件转换
    - 智能标注模式：
      * 特定路径模式：检测到handwriting/printer目录时，使用整个图片作为标注
      * PaddleX模式：使用PaddleX进行layout detection
    - 自动生成训练/验证数据分割
    - 输出符合COCO格式的标注文件

使用方法:
    # 使用默认参数
    python continue_gen_dataset.py --input_path data/source --output_path data/dataset
    
    # 指定训练比例和模型
    python continue_gen_dataset.py --input_path data/source --output_path data/dataset --train_ratio 0.8 --model_name PP-DocLayout-L --model_dir output/best_model/inference
    
    # 查看帮助信息
    python continue_gen_dataset.py --help

参数说明:
    --input_path: 源文件路径，包含需要处理的文件
    --output_path: 输出数据集路径
    --train_ratio: 训练数据比例，默认0.8（80%训练，20%验证）
    --model_name: 版面检测模型，可选PP-DocLayout_plus-L或PP-DocLayout-L

输出格式:
    生成的数据集包含以下结构:
    - images/: 所有图片文件
    - annotations/: 标注文件
      - instance_train.json: 训练数据标注
      - instance_val.json: 验证数据标注
    - split_jsons/: 每个图片的单独标注文件

注意事项:
    - 需要安装PaddleX: pip install paddlex
    - 需要安装文档处理库: pip install python-docx pdf2image
    - 建议使用conda环境: conda activate seal_env
"""

import os
import json
import argparse
import shutil
import glob
from typing import Dict, List, Any, Tuple
import subprocess
import sys
from pathlib import Path

# 尝试导入必要的库
PIL_AVAILABLE = False
PADDLEOCR_AVAILABLE = False
DOCX_AVAILABLE = False
PDF_AVAILABLE = False

try:
    from PIL import Image
    PIL_AVAILABLE = True
    print("✓ PIL/Pillow库可用")
except ImportError:
    print("⚠ 警告: PIL库未安装，请运行: pip install Pillow")

try:
    from paddlex import create_model
    PADDLEX_AVAILABLE = True
    print("✓ PaddleX库可用")
except ImportError:
    PADDLEX_AVAILABLE = False
    print("⚠ 警告: PaddleX库未安装，请运行: pip install paddlex")

try:
    from docx import Document
    DOCX_AVAILABLE = True
    print("✓ python-docx库可用")
except ImportError:
    print("⚠ 警告: python-docx库未安装，请运行: pip install python-docx")

try:
    from pdf2image import convert_from_path
    PDF_AVAILABLE = True
    print("✓ pdf2image库可用")
except ImportError:
    print("⚠ 警告: pdf2image库未安装，请运行: pip install pdf2image")


class AutoDatasetGenerator:
    """自动化数据集生成器"""
    
    def __init__(self, train_ratio: float = 0.8, model_name: str = "PP-DocLayout_plus-L", model_dir: str = None):
        self.train_ratio = train_ratio
        self.model_name = model_name
        self.model_dir = model_dir
        self.ocr = None
        self.layout_model = None
        self.coco_data = {
            "info": {
                "description": "Auto Generated Dataset with PaddleOCR",
                "version": "1.0",
                "year": 2024,
                "contributor": "Auto Dataset Generator",
                "date_created": "2025-01-20"
            },
            "licenses": [
                {
                    "id": 0,
                    "name": "Unknown",
                    "url": None
                }
            ],
            "images": [],
            "annotations": [],
            "categories": [
                {
                    "id": 0,
                    "name": "_background_",
                    "supercategory": None
                },
                {
                    "id": 1,
                    "name": "Text",
                    "supercategory": None
                },
                {
                    "id": 2,
                    "name": "Title",
                    "supercategory": None
                },
                {
                    "id": 3,
                    "name": "Figure",
                    "supercategory": None
                },
                {
                    "id": 4,
                    "name": "Figure caption",
                    "supercategory": None
                },
                {
                    "id": 5,
                    "name": "Table",
                    "supercategory": None
                },
                {
                    "id": 6,
                    "name": "Table caption",
                    "supercategory": None
                },
                {
                    "id": 7,
                    "name": "Header",
                    "supercategory": None
                },
                {
                    "id": 8,
                    "name": "Footer",
                    "supercategory": None
                },
                {
                    "id": 9,
                    "name": "Seal",
                    "supercategory": None
                },
                {
                    "id": 10,
                    "name": "Equation",
                    "supercategory": None
                },
                {
                    "id": 11,
                    "name": "Handwriting",
                    "supercategory": None
                }
            ],
            "type": "instances"
        }
        self.image_id = 0
        self.annotation_id = 0
        self.all_images = []
        self.all_annotations = []
        self.existing_data = {
            'train_images': [],
            'train_annotations': [],
            'val_images': [],
            'val_annotations': []
        }
    
    def initialize_ocr(self):
        """初始化PaddleOCR（使用Python API）"""
        try:
            print("正在初始化PaddleOCR...")
            
            # 尝试导入PaddleOCR
            try:
                from paddleocr import PaddleOCR
            except ImportError:
                print("✗ PaddleOCR未安装，请运行: pip install paddleocr")
                return False
            
            # 尝试不同的初始化参数
            init_methods = [
                # 方法1: 最新参数
                {
                    "params": {"use_textline_orientation": True, "lang": "ch", "show_log": False},
                    "name": "最新参数"
                },
                # 方法2: 旧参数
                {
                    "params": {"use_angle_cls": True, "lang": "ch", "show_log": False},
                    "name": "旧参数"
                },
                # 方法3: 简化参数
                {
                    "params": {"lang": "ch", "show_log": False},
                    "name": "简化参数"
                },
                # 方法4: 最简参数
                {
                    "params": {"show_log": False},
                    "name": "最简参数"
                },
                # 方法5: 无参数
                {
                    "params": {},
                    "name": "无参数"
                }
            ]
            
            for method in init_methods:
                try:
                    print(f"  尝试{method['name']}...")
                    self.ocr = PaddleOCR(**method['params'])
                    print(f"✓ PaddleOCR初始化成功 ({method['name']})")
                    return True
                except Exception as e:
                    print(f"  ✗ {method['name']}失败: {e}")
                    continue
            
            print("✗ 所有初始化方法都失败")
            return False
                
        except Exception as e:
            print(f"✗ PaddleOCR初始化失败: {e}")
            return False
    
    def initialize_layout_model(self):
        """初始化PaddleX版面检测模型"""
        try:
            print(f"正在初始化PaddleX版面检测模型: {self.model_name}")
            
            if not PADDLEX_AVAILABLE:
                print("✗ PaddleX未安装，请运行: pip install paddlex")
                return False
            
            try:
                import paddle
                print(f"✓ PaddlePaddle: {paddle.__version__}")
            except Exception as e:
                print(f"✗ PaddlePaddle导入失败: {e}")
                return False
            
            try:
                self.layout_model = create_model(model_name=self.model_name,model_dir=self.model_dir)
                print(f"✓ PaddleX模型初始化成功: {self.model_name} {self.model_dir}")
                return True
            except Exception as e:
                print(f"✗ PaddleX模型加载失败: {e}")
                return False
                
        except Exception as e:
            print(f"✗ PaddleX模型初始化失败: {e}")
            return False
    
    def get_image_size(self, image_path: str) -> Tuple[int, int]:
        """获取图片尺寸"""
        try:
            if PIL_AVAILABLE:
                with Image.open(image_path) as img:
                    return img.size  # (width, height)
            else:
                # 使用内置方法
                return self._get_image_size_from_file(image_path)
        except Exception as e:
            print(f"错误: 无法获取图片尺寸 {image_path}: {e}")
            return (800, 600)  # 默认尺寸
    
    def _get_image_size_from_file(self, image_path: str) -> Tuple[int, int]:
        """使用内置方法获取图片尺寸"""
        try:
            import struct
            with open(image_path, 'rb') as f:
                header = f.read(24)
                
                # PNG格式
                if header[:8] == b'\x89PNG\r\n\x1a\n':
                    width = struct.unpack('>I', header[16:20])[0]
                    height = struct.unpack('>I', header[20:24])[0]
                    return width, height
                
                # JPEG格式
                elif header[:2] == b'\xff\xd8':
                    f.seek(0)
                    while True:
                        marker = f.read(2)
                        if not marker or marker[0] != 0xff:
                            break
                        if marker[1] == 0xd8:  # SOI
                            continue
                        if marker[1] == 0xd9:  # EOI
                            break
                        if marker[1] == 0xda:  # SOS
                            break
                        length = struct.unpack('>H', f.read(2))[0]
                        if marker[1] in [0xc0, 0xc1, 0xc2]:  # SOF0, SOF1, SOF2
                            f.read(1)  # precision
                            height = struct.unpack('>H', f.read(2))[0]
                            width = struct.unpack('>H', f.read(2))[0]
                            return width, height
                        f.read(length - 2)
                
                return (800, 600)  # 默认尺寸
        except Exception as e:
            print(f"错误: 无法读取图片尺寸 {image_path}: {e}")
            return (800, 600)
    
    def _delete_related_annotations(self, filename: str, output_path: str):
        """删除与指定文件相关的标注数据"""
        try:
            # 删除split_jsons目录中的对应JSON文件
            base_name = os.path.splitext(filename)[0]
            split_jsons_dir = os.path.join(output_path, "split_jsons")
            
            if os.path.exists(split_jsons_dir):
                json_file = os.path.join(split_jsons_dir, f"{base_name}.json")
                if os.path.exists(json_file):
                    os.remove(json_file)
                    print(f"✓ 已删除相关标注文件: {base_name}.json")
            
            # 从现有的汇总数据中删除相关记录
            self._remove_from_existing_data(filename)
            
        except Exception as e:
            print(f"✗ 删除相关标注数据失败 {filename}: {e}")
    
    def _remove_from_existing_data(self, filename: str):
        """从现有数据中移除指定文件的相关记录"""
        try:
            # 查找并移除对应的图片记录
            image_to_remove = None
            for image in self.existing_data['images']:
                if image['file_name'] == filename:
                    image_to_remove = image
                    break
            
            if image_to_remove:
                # 移除图片记录
                self.existing_data['images'].remove(image_to_remove)
                print(f"✓ 已从现有数据中移除图片记录: {filename}")
                
                # 移除相关的标注记录
                annotations_to_remove = []
                for annotation in self.existing_data['annotations']:
                    if annotation['image_id'] == image_to_remove['id']:
                        annotations_to_remove.append(annotation)
                
                for annotation in annotations_to_remove:
                    self.existing_data['annotations'].remove(annotation)
                
                if annotations_to_remove:
                    print(f"✓ 已从现有数据中移除 {len(annotations_to_remove)} 个相关标注记录")
                    
        except Exception as e:
            print(f"✗ 从现有数据中移除记录失败 {filename}: {e}")
    
    def _load_existing_data(self, output_path: str):
        """加载现有的数据集"""
        print("检查现有数据集...")
        
        # 检查是否存在现有的标注文件
        train_file = os.path.join(output_path, "annotations", "instance_train.json")
        val_file = os.path.join(output_path, "annotations", "instance_val.json")
        
        existing_train_images = []
        existing_train_annotations = []
        existing_val_images = []
        existing_val_annotations = []
        
        # 加载训练数据
        if os.path.exists(train_file):
            try:
                with open(train_file, 'r', encoding='utf-8') as f:
                    train_data = json.load(f)
                    existing_train_images.extend(train_data.get('images', []))
                    existing_train_annotations.extend(train_data.get('annotations', []))
                print(f"✓ 加载现有训练数据: {len(train_data.get('images', []))} 个图片")
            except Exception as e:
                print(f"⚠ 加载训练数据失败: {e}")
        
        # 加载验证数据
        if os.path.exists(val_file):
            try:
                with open(val_file, 'r', encoding='utf-8') as f:
                    val_data = json.load(f)
                    existing_val_images.extend(val_data.get('images', []))
                    existing_val_annotations.extend(val_data.get('annotations', []))
                print(f"✓ 加载现有验证数据: {len(val_data.get('images', []))} 个图片")
            except Exception as e:
                print(f"⚠ 加载验证数据失败: {e}")
        
        # 重新分配ID以避免冲突
        all_existing_images = existing_train_images + existing_val_images
        all_existing_annotations = existing_train_annotations + existing_val_annotations
        
        max_image_id = max([img['id'] for img in all_existing_images], default=-1)
        max_annotation_id = max([ann['id'] for ann in all_existing_annotations], default=-1)
        
        self.image_id = max_image_id + 1
        self.annotation_id = max_annotation_id + 1
        
        self.existing_data['train_images'] = existing_train_images
        self.existing_data['train_annotations'] = existing_train_annotations
        self.existing_data['val_images'] = existing_val_images
        self.existing_data['val_annotations'] = existing_val_annotations
        
        print(f"✓ 现有训练数据: {len(existing_train_images)} 个图片和 {len(existing_train_annotations)} 个标注")
        print(f"✓ 现有验证数据: {len(existing_val_images)} 个图片和 {len(existing_val_annotations)} 个标注")
        print(f"✓ 新数据将从ID {self.image_id} 开始")
    
    def copy_image_files(self, input_path: str, output_path: str) -> List[str]:
        """拷贝图片文件到输出目录，处理重名文件"""
        print(f"开始拷贝图片文件从 {input_path} 到 {output_path}")
        
        # 创建输出目录
        os.makedirs(output_path, exist_ok=True)
        
        # 支持的图片格式
        image_extensions = ['*.png', '*.jpg', '*.jpeg', '*.bmp', '*.tiff', '*.tif', '*.gif', '*.webp']
        
        copied_files = []
        for ext in image_extensions:
            pattern = os.path.join(input_path, '**', ext)
            for file_path in glob.glob(pattern, recursive=True):
                filename = os.path.basename(file_path)
                dest_path = os.path.join(output_path, filename)
                
                # 检查目标文件是否已存在
                if os.path.exists(dest_path):
                    print(f"⚠ 警告: 发现重名文件 {filename}，将删除现有文件和相关标注数据")
                    
                    # 删除现有的JSON标注文件
                    self._delete_related_annotations(filename, output_path)
                    
                    # 删除现有文件
                    try:
                        os.remove(dest_path)
                        print(f"✓ 已删除现有文件: {filename}")
                    except Exception as e:
                        print(f"✗ 删除现有文件失败 {dest_path}: {e}")
                
                try:
                    shutil.copy2(file_path, dest_path)
                    copied_files.append(dest_path)
                    print(f"✓ 拷贝图片: {filename}")
                except Exception as e:
                    print(f"✗ 拷贝失败 {file_path}: {e}")
        
        print(f"共拷贝了 {len(copied_files)} 个图片文件")
        return copied_files
    
    def convert_docx_to_images(self, docx_path: str, output_path: str) -> List[str]:
        """将Word文档转换为图片"""
        if not DOCX_AVAILABLE:
            print("⚠ python-docx库未安装，跳过Word文档转换")
            return []
        
        print(f"转换Word文档: {docx_path}")
        converted_files = []
        
        try:
            # 这里需要实现Word到图片的转换
            # 由于python-docx主要处理文本，图片转换需要其他方法
            # 可以使用python-docx2txt + 截图的方式，或者使用其他库
            print("⚠ Word文档转换功能需要额外实现")
            return []
        except Exception as e:
            print(f"✗ Word转换失败 {docx_path}: {e}")
            return []
    
    def convert_pdf_to_images(self, pdf_path: str, output_path: str) -> List[str]:
        """将PDF文件转换为图片，处理重名文件"""
        if not PDF_AVAILABLE:
            print("⚠ pdf2image库未安装，跳过PDF转换")
            return []
        
        print(f"转换PDF文件: {pdf_path}")
        converted_files = []
        
        try:
            # 获取PDF文件名（不含扩展名）
            pdf_name = os.path.splitext(os.path.basename(pdf_path))[0]
            
            # 检查poppler是否安装
            if not self._check_poppler_installation():
                print("⚠ poppler未安装，跳过PDF转换")
                print("请安装poppler:")
                print("  macOS: brew install poppler")
                print("  Linux: sudo apt-get install poppler-utils")
                print("  Windows: 下载poppler并添加到PATH")
                return []
            
            # 转换PDF为图片，添加错误处理
            try:
                images = convert_from_path(pdf_path, dpi=300, first_page=1, last_page=None)
            except Exception as e:
                print(f"✗ PDF转换失败: {e}")
                print("尝试使用不同的参数...")
                try:
                    # 尝试降低DPI
                    images = convert_from_path(pdf_path, dpi=150, first_page=1, last_page=None)
                except Exception as e2:
                    print(f"✗ 降低DPI后仍然失败: {e2}")
                    return []
            
            for i, image in enumerate(images):
                # 按 文件名_页码.jpg 格式命名
                output_filename = f"{pdf_name}_{i+1}.jpg"
                output_filepath = os.path.join(output_path, output_filename)
                
                # 检查目标文件是否已存在
                if os.path.exists(output_filepath):
                    print(f"⚠ 警告: 发现重名文件 {output_filename}，将删除现有文件和相关标注数据")
                    
                    # 删除现有的JSON标注文件
                    self._delete_related_annotations(output_filename, output_path)
                    
                    # 删除现有文件
                    try:
                        os.remove(output_filepath)
                        print(f"✓ 已删除现有文件: {output_filename}")
                    except Exception as e:
                        print(f"✗ 删除现有文件失败 {output_filepath}: {e}")
                
                # 保存图片
                image.save(output_filepath, 'JPEG', quality=95)
                converted_files.append(output_filepath)
                print(f"✓ 转换页面 {i+1}: {output_filename}")
            
            print(f"✓ PDF转换完成，共 {len(converted_files)} 页")
            return converted_files
            
        except Exception as e:
            print(f"✗ PDF转换失败 {pdf_path}: {e}")
            return []
    
    def _check_poppler_installation(self) -> bool:
        """检查poppler是否安装"""
        try:
            import subprocess
            # 检查pdftoppm命令是否可用
            result = subprocess.run(['pdftoppm', '-h'], capture_output=True, text=True, timeout=5)
            return result.returncode == 0
        except (subprocess.TimeoutExpired, FileNotFoundError):
            return False
        except Exception:
            return False
    
    def process_documents(self, input_path: str, output_path: str) -> List[str]:
        """处理文档文件（Word、PDF）转换为图片"""
        print(f"开始处理文档文件从 {input_path}")
        
        # 创建输出目录
        os.makedirs(output_path, exist_ok=True)
        
        all_converted_files = []
        
        # 处理PDF文件
        pdf_pattern = os.path.join(input_path, '**', '*.pdf')
        for pdf_path in glob.glob(pdf_pattern, recursive=True):
            converted_files = self.convert_pdf_to_images(pdf_path, output_path)
            all_converted_files.extend(converted_files)
        
        # 处理Word文件
        docx_pattern = os.path.join(input_path, '**', '*.docx')
        for docx_path in glob.glob(docx_pattern, recursive=True):
            converted_files = self.convert_docx_to_images(docx_path, output_path)
            all_converted_files.extend(converted_files)
        
        print(f"文档转换完成，共生成 {len(all_converted_files)} 个图片文件")
        return all_converted_files
    
    def run_paddleocr_detection(self, image_path: str) -> List[Dict]:
        """使用PaddleOCR进行目标检测（兼容3.2.0版本）"""
        if not self.ocr:
            print("错误: PaddleOCR未初始化")
            return []
        
        try:
            print(f"正在检测图片: {os.path.basename(image_path)}")
            
            # 使用Python API调用PaddleOCR
            result = None
            
            # 尝试不同的API调用方式，添加内存管理
            try:
                import gc
                gc.collect()  # 强制垃圾回收
                
                # 优先使用predict方法（PaddleOCR 3.2.0新版本API）
                if hasattr(self.ocr, 'predict'):
                    try:
                        result = self.ocr.predict(image_path)
                        print("✓ 使用predict方法成功")
                    except Exception as e:
                        print(f"predict方法失败: {e}")
                        # 回退到ocr方法
                        if hasattr(self.ocr, 'ocr'):
                            result = self.ocr.ocr(image_path)
                            print("✓ 使用ocr方法成功")
                        else:
                            print("✗ PaddleOCR对象没有可用的方法")
                            return []
                elif hasattr(self.ocr, 'ocr'):
                    result = self.ocr.ocr(image_path)
                    print("✓ 使用ocr方法成功")
                else:
                    print("✗ PaddleOCR对象没有ocr或predict方法")
                    return []
                
                # 再次垃圾回收
                gc.collect()
                
            except Exception as e:
                print(f"✗ PaddleOCR API调用失败: {e}")
                # 尝试重新初始化OCR
                try:
                    print("尝试重新初始化PaddleOCR...")
                    self.ocr = None
                    if self.initialize_ocr():
                        # 重试一次
                        if hasattr(self.ocr, 'predict'):
                            result = self.ocr.predict(image_path)
                        elif hasattr(self.ocr, 'ocr'):
                            result = self.ocr.ocr(image_path)
                    else:
                        return []
                except Exception as retry_error:
                    print(f"重新初始化失败: {retry_error}")
                    return []
            
            detections = []
            if result and result[0]:
                for line in result[0]:
                    if line:
                        # 解析检测结果
                        bbox = line[0]  # 边界框坐标
                        text_info = line[1]  # (文本, 置信度)
                        
                        # 转换为COCO格式的边界框 [x, y, width, height]
                        x_coords = [point[0] for point in bbox]
                        y_coords = [point[1] for point in bbox]
                        x_min, x_max = min(x_coords), max(x_coords)
                        y_min, y_max = min(y_coords), max(y_coords)
                        width = x_max - x_min
                        height = y_max - y_min
                        
                        # 根据文本内容智能分类
                        category_id = self._classify_text_category(text_info[0] if text_info else '')
                        
                        detection = {
                            'bbox': [x_min, y_min, width, height],
                            'text': text_info[0] if text_info else '',
                            'confidence': text_info[1] if text_info else 0.0,
                            'category_id': category_id
                        }
                        detections.append(detection)
            
            print(f"✓ 检测到 {len(detections)} 个目标")
            return detections
            
        except Exception as e:
            print(f"✗ PaddleOCR检测失败 {image_path}: {e}")
            return []
    
    def run_paddlex_layout_detection(self, image_path: str) -> List[Dict]:
        """使用PaddleX进行版面检测"""
        if not self.layout_model:
            print("错误: PaddleX版面检测模型未初始化")
            return []
        
        try:
            print(f"正在使用PaddleX检测图片: {os.path.basename(image_path)}")
            
            # 使用PaddleX进行版面检测
            output = self.layout_model.predict(image_path, batch_size=1, layout_nms=True)
            
            # 将生成器转换为列表
            output_list = list(output)
            
            detections = []
            if output_list and isinstance(output_list, list):
                for res in output_list:
                    # 检查结果结构 - 通过字典方式访问boxes
                    if 'boxes' in res and res['boxes']:
                        boxes = res['boxes']
                        for bbox_info in boxes:
                            # 解析检测结果
                            coordinate = bbox_info['coordinate']  # [x_min, y_min, x_max, y_max]
                            label = bbox_info['label']
                            score = bbox_info['score']
                            
                            # 转换为COCO格式 [x_min, y_min, width, height]
                            x_min, y_min, x_max, y_max = coordinate
                            width = x_max - x_min
                            height = y_max - y_min
                            # 转换为Python原生类型以避免JSON序列化错误
                            bbox = [float(x_min), float(y_min), float(width), float(height)]
                            
                            # 根据标签映射到COCO类别
                            category_id = self._map_paddlex_label_to_coco(label)
                            
                            detection = {
                                'bbox': bbox,
                                'text': label,
                                'confidence': float(score),  # 转换为Python原生类型
                                'category_id': category_id
                            }
                            detections.append(detection)
            
            print(f"✓ PaddleX检测到 {len(detections)} 个目标")
            return detections
            
        except Exception as e:
            print(f"✗ PaddleX版面检测失败 {image_path}: {e}")
            return []
    
    def _map_paddlex_label_to_coco(self, label: str) -> int:
        """将PaddleX标签映射到COCO类别ID"""
        label_mapping = {
            'text': 1,                    # Text
            'title': 2,                   # Title
            'paragraph_title': 2,         # Title
            'figure': 3,                  # Figure
            'figure_caption': 4,          # Figure caption
            'table': 5,                   # Table
            'table_caption': 6,           # Table caption
            'header': 7,                  # Header
            'footer': 8,                  # Footer
            'reference': 1,               # Reference
            'equation': 10,               # Equation
            'handwriting': 11,            # Handwriting
            'content': 1,                 # 默认为Text
            'seal': 9,                    # 印章归类为Seal
            'image': 3,                   # 图片归类为Figure
        }
        
        return label_mapping.get(label.lower(), 1)  # 默认为Text类别
    
    def _check_specific_path_structure(self, input_path: str) -> bool:
        """检查输入路径是否具有特定的目录结构（handwriting/printer）"""
        handwriting_path = os.path.join(input_path, "handwriting")
        printer_path = os.path.join(input_path, "printer")
        
        # 检查是否存在handwriting或printer子目录
        has_handwriting = os.path.exists(handwriting_path) and os.path.isdir(handwriting_path)
        has_printer = os.path.exists(printer_path) and os.path.isdir(printer_path)
        
        if has_handwriting or has_printer:
            print(f"✓ 检测到特定目录结构:")
            if has_handwriting:
                print(f"  - handwriting目录: {handwriting_path}")
            if has_printer:
                print(f"  - printer目录: {printer_path}")
            return True
        
        return False
    
    def _process_specific_path_annotation(self, input_path: str, output_path: str):
        """处理特定路径结构的标注（基于generate_json_annotations.py的逻辑）"""
        print("=" * 60)
        print("使用特定路径标注逻辑")
        print("=" * 60)
        
        # 创建输出目录
        images_dir = os.path.join(output_path, "images")
        split_jsons_dir = os.path.join(output_path, "split_jsons")
        os.makedirs(images_dir, exist_ok=True)
        os.makedirs(split_jsons_dir, exist_ok=True)
        
        total_files = 0
        
        # 处理handwriting文件夹 - 手写体标注为handwriting (category_id=11)
        handwriting_path = os.path.join(input_path, "handwriting")
        if os.path.exists(handwriting_path):
            print(f"处理handwriting文件夹: {handwriting_path}")
            count = self._process_folder_with_specific_annotation(
                handwriting_path, images_dir, split_jsons_dir, category_id=11, label="handwriting"
            )
            total_files += count
        else:
            print(f"警告: handwriting文件夹不存在: {handwriting_path}")
        
        # 处理printer文件夹 - 印刷体标注为text (category_id=1)
        printer_path = os.path.join(input_path, "printer")
        if os.path.exists(printer_path):
            print(f"处理printer文件夹: {printer_path}")
            count = self._process_folder_with_specific_annotation(
                printer_path, images_dir, split_jsons_dir, category_id=1, label="text"
            )
            total_files += count
        else:
            print(f"警告: printer文件夹不存在: {printer_path}")
        
        print(f"特定路径标注完成，共处理 {total_files} 个图片文件")
        
        # 生成训练和验证数据分割
        self.generate_train_val_split(output_path)
    
    def _process_folder_with_specific_annotation(self, folder_path: str, images_dir: str, 
                                               split_jsons_dir: str, category_id: int, label: str) -> int:
        """处理单个文件夹中的所有图片，使用特定标注逻辑"""
        # 支持的图片格式
        image_extensions = ['*.png', '*.jpg', '*.jpeg', '*.bmp', '*.tiff', '*.tif']
        
        image_files = []
        for ext in image_extensions:
            pattern = os.path.join(folder_path, ext)
            image_files.extend(glob.glob(pattern))
        
        print(f"在{label}文件夹中找到 {len(image_files)} 个图片文件")
        
        processed_count = 0
        for image_path in image_files:
            # 获取图片信息
            width, height = self.get_image_size(image_path)
            image_info = {
                "id": self.image_id,
                "file_name": os.path.basename(image_path),
                "width": width,
                "height": height
            }
            
            # 生成新的文件名（添加路径前缀）
            original_filename = os.path.basename(image_path)
            name, ext = os.path.splitext(original_filename)
            new_filename = f"{label}_{name}{ext}"
            new_image_path = os.path.join(images_dir, new_filename)
            
            # 复制图片到统一目录
            try:
                shutil.copy2(image_path, new_image_path)
                print(f"复制图片: {original_filename} -> {new_filename}")
            except Exception as e:
                print(f"复制图片失败 {image_path}: {e}")
                continue
            
            # 更新图片信息中的文件名
            image_info["file_name"] = new_filename
            
            # 创建特定标注（使用整个图片作为边界框）
            annotation = self._create_specific_annotation(image_info, category_id, label)
            
            # 构建单个图片的COCO数据
            single_coco_data = {
                "info": self.coco_data["info"].copy(),
                "licenses": self.coco_data["licenses"].copy(),
                "images": [image_info],
                "annotations": [annotation],
                "categories": self.coco_data["categories"].copy(),
                "type": self.coco_data["type"]
            }
            
            # 将数据添加到汇总列表中
            self.all_images.append(image_info)
            self.all_annotations.append(annotation)
            
            # 生成输出文件名
            base_name = os.path.splitext(original_filename)[0]
            output_file = os.path.join(split_jsons_dir, f"{label}_{base_name}.json")
            
            # 保存JSON文件
            with open(output_file, 'w', encoding='utf-8') as f:
                json.dump(single_coco_data, f, ensure_ascii=False, indent=2)
            
            processed_count += 1
            self.image_id += 1
            self.annotation_id += 1
            
            # 每处理100个文件打印一次进度
            if processed_count % 100 == 0:
                print(f"已处理 {processed_count} 个{label}图片文件")
        
        return processed_count
    
    def _create_specific_annotation(self, image_info: Dict[str, Any], category_id: int, label: str) -> Dict[str, Any]:
        """创建特定标注信息（使用整个图片作为边界框）"""
        # 使用整个图片作为边界框
        bbox = [0, 0, image_info["width"], image_info["height"]]
        area = image_info["width"] * image_info["height"]
        
        # 创建完整坐标的segmentation（矩形四个角点）
        width = image_info["width"]
        height = image_info["height"]
        segmentation = [
            [0, 0, width, 0, width, height, 0, height]  # 左上、右上、右下、左下
        ]
        
        annotation = {
            "id": self.annotation_id,
            "image_id": image_info["id"],
            "category_id": category_id,
            "bbox": bbox,
            "area": area,
            "iscrowd": 0,
            "segmentation": segmentation
        }
        
        return annotation
    
    def _run_paddleocr_command(self, image_path: str) -> List:
        """使用命令行调用PaddleOCR"""
        try:
            import subprocess
            import json
            import tempfile
            import os
            
            # 构建PaddleOCR命令 - 使用layout_detection，不指定输出文件
            cmd = [
                'paddleocr', 
                'layout_detection',
                '-i', image_path
            ]
            
            # 执行命令
            result = subprocess.run(cmd, capture_output=True, text=True, timeout=120)
            
            if result.returncode != 0:
                print(f"PaddleOCR命令执行失败: {result.stderr}")
                return []
            
            # 解析输出结果
            output_text = result.stdout.strip()
            if not output_text:
                print("PaddleOCR未返回结果")
                return []
            
            # 尝试解析JSON输出
            try:
                # 查找JSON部分（通常在输出的最后）
                lines = output_text.split('\n')
                json_start = -1
                for i, line in enumerate(lines):
                    if line.strip().startswith('{') or line.strip().startswith('['):
                        json_start = i
                        break
                
                if json_start >= 0:
                    json_text = '\n'.join(lines[json_start:])
                    data = json.loads(json_text)
                    return data.get('results', [])
                else:
                    # 如果没有找到JSON，尝试解析文本输出
                    return self._parse_text_output(output_text)
                    
            except json.JSONDecodeError:
                # 如果JSON解析失败，尝试解析文本输出
                return self._parse_text_output(output_text)
            
        except subprocess.TimeoutExpired:
            print("PaddleOCR命令超时")
            return []
        except Exception as e:
            print(f"命令行调用PaddleOCR失败: {e}")
            return []
    
    def _parse_text_output(self, output_text: str) -> List:
        """解析PaddleOCR的文本输出"""
        try:
            # 这是一个简化的解析器，根据实际输出格式调整
            lines = output_text.split('\n')
            results = []
            
            for line in lines:
                line = line.strip()
                if line and not line.startswith('[') and not line.startswith('{'):
                    # 假设每行包含检测结果
                    # 这里需要根据实际的PaddleOCR输出格式来调整
                    # 暂时返回空结果，避免错误
                    pass
            
            return results
            
        except Exception as e:
            print(f"解析PaddleOCR输出失败: {e}")
            return []
    
    def _classify_text_category(self, text: str) -> int:
        """根据文本内容智能分类"""
        if not text:
            return 1  # 默认为Text
        
        text_lower = text.lower().strip()
        
        # 标题特征
        if (len(text) < 50 and 
            (text_lower.startswith(('title', '标题', '第', 'chapter', 'section')) or
             text.isupper() or
             any(keyword in text_lower for keyword in ['标题', '题目', '名称', 'name', 'title']))):
            return 2  # Title
        
        # 表格特征
        if (any(keyword in text_lower for keyword in ['表', 'table', '表格', '数据', 'data']) or
            '|' in text or
            '\t' in text):
            return 5  # Table
        
        # 图片/图表特征
        if any(keyword in text_lower for keyword in ['图', 'figure', 'fig', '图片', 'chart', '图表']):
            return 3  # Figure
        
        # 公式特征
        if any(keyword in text_lower for keyword in ['公式', 'equation', 'eq', '=', '∑', '∫', '√']):
            return 10  # Equation
        
        # 页眉页脚特征
        if (text_lower in ['header', 'footer', '页眉', '页脚'] or
            any(keyword in text_lower for keyword in ['page', '页', '页码'])):
            return 7 if 'header' in text_lower or '页眉' in text_lower else 8  # Header/Footer
        
        # 手写体特征（简单判断）
        if any(keyword in text_lower for keyword in ['手写', 'handwriting', '签名', 'signature']):
            return 11  # Handwriting
        
        # 印章特征
        if any(keyword in text_lower for keyword in ['章', '印', '公司', '银行', '财务专用章', '公章', '法人章', '合同专用章', '发票专用章', '财务专用章', '公章', '法人章', '合同专用章', '发票专用章']):
            return 9  # Seal
        
        # 参考文献特征
        if any(keyword in text_lower for keyword in ['ref', '参考文献', 'reference', '引用', 'cite']):
            return 1  # Reference
        
        # 默认为普通文本
        return 1  # Text
    
    def create_annotation_from_detection(self, image_info: Dict, detection: Dict) -> Dict:
        """从检测结果创建标注"""
        annotation = {
            "id": self.annotation_id,
            "image_id": image_info["id"],
            "category_id": detection["category_id"],
            "bbox": detection["bbox"],
            "area": detection["bbox"][2] * detection["bbox"][3],
            "iscrowd": 0,
            "segmentation": [
                [
                    detection["bbox"][0], detection["bbox"][1],
                    detection["bbox"][0] + detection["bbox"][2], detection["bbox"][1],
                    detection["bbox"][0] + detection["bbox"][2], detection["bbox"][1] + detection["bbox"][3],
                    detection["bbox"][0], detection["bbox"][1] + detection["bbox"][3]
                ]
            ]
        }
        
        self.annotation_id += 1
        return annotation
    
    def process_images_with_ocr(self, image_files: List[str], output_path: str):
        """使用PaddleX处理所有图片并生成标注"""
        print(f"开始使用PaddleX处理 {len(image_files)} 个图片文件")
        
        # 创建输出目录
        annotations_dir = os.path.join(output_path, "annotations")
        split_jsons_dir = os.path.join(output_path, "split_jsons")
        os.makedirs(annotations_dir, exist_ok=True)
        os.makedirs(split_jsons_dir, exist_ok=True)
        
        processed_count = 0
        failed_count = 0
        
        for i, image_path in enumerate(image_files):
            try:
                # 获取图片信息
                width, height = self.get_image_size(image_path)
                image_info = {
                    "id": self.image_id,
                    "file_name": os.path.basename(image_path),
                    "width": width,
                    "height": height
                }
                
                # 使用PaddleX进行版面检测
                detections = self.run_paddlex_layout_detection(image_path)
                
                # 创建标注
                annotations = []
                for detection in detections:
                    annotation = self.create_annotation_from_detection(image_info, detection)
                    annotations.append(annotation)
                
                # 构建单个图片的COCO数据
                single_coco_data = {
                    "info": self.coco_data["info"].copy(),
                    "licenses": self.coco_data["licenses"].copy(),
                    "images": [image_info],
                    "annotations": annotations,
                    "categories": self.coco_data["categories"].copy(),
                    "type": self.coco_data["type"]
                }
                
                # 保存单个图片的标注文件
                base_name = os.path.splitext(os.path.basename(image_path))[0]
                split_json_file = os.path.join(split_jsons_dir, f"{base_name}.json")
                with open(split_json_file, 'w', encoding='utf-8') as f:
                    json.dump(single_coco_data, f, ensure_ascii=False, indent=2)
                
                # 添加到汇总数据
                self.all_images.append(image_info)
                self.all_annotations.extend(annotations)
                
                self.image_id += 1
                processed_count += 1
                
                print(f"✓ 处理完成: {os.path.basename(image_path)} ({len(annotations)} 个标注)")
                
                # 每处理10个文件后强制垃圾回收
                if processed_count % 10 == 0:
                    import gc
                    gc.collect()
                    print(f"  已处理 {processed_count}/{len(image_files)} 个文件，进行内存清理...")
                
            except Exception as e:
                failed_count += 1
                print(f"✗ 处理失败: {os.path.basename(image_path)} - {e}")
                
                # 如果连续失败太多，尝试重新初始化OCR
                if failed_count > 5:
                    print("⚠ 连续失败过多，尝试重新初始化PaddleOCR...")
                    try:
                        self.ocr = None
                        if not self.initialize_ocr():
                            print("✗ 重新初始化失败，跳过剩余文件")
                            break
                        failed_count = 0
                    except Exception as reinit_error:
                        print(f"✗ 重新初始化失败: {reinit_error}")
                        break
        
        print(f"所有图片处理完成，共生成 {len(self.all_images)} 个图片的标注")
        if failed_count > 0:
            print(f"⚠ 处理失败: {failed_count} 个文件")
    
    def generate_train_val_split(self, output_path: str):
        """生成训练和验证数据分割，保持已有数据不变，只对新数据进行划分"""
        print("开始生成训练和验证数据分割...")
        
        # 保持已有数据不变
        existing_train_images = self.existing_data['train_images']
        existing_train_annotations = self.existing_data['train_annotations']
        existing_val_images = self.existing_data['val_images']
        existing_val_annotations = self.existing_data['val_annotations']
        
        print(f"✓ 保持已有训练数据不变: {len(existing_train_images)} 个图片")
        print(f"✓ 保持已有验证数据不变: {len(existing_val_images)} 个图片")
        print(f"✓ 新增数据: {len(self.all_images)} 个图片")
        
        # 对新数据进行train_ratio划分
        new_data_count = len(self.all_images)
        if new_data_count > 0:
            # 限制train_ratio在合理范围内
            if self.train_ratio >= 1.0:
                new_train_count = new_data_count  # 所有新数据进入训练集
                new_val_count = 0  # 没有新数据进入验证集
            else:
                new_train_count = int(new_data_count * self.train_ratio)
                new_val_count = new_data_count - new_train_count
            
            print(f"✓ 新数据划分: {new_train_count} 个训练 + {new_val_count} 个验证")
            
            # 分割新数据
            new_train_images = self.all_images[:new_train_count]
            new_val_images = self.all_images[new_train_count:]
            
            # 分割新数据的标注
            new_train_annotations = []
            new_val_annotations = []
            
            # 根据新数据的image_id范围来分割标注
            new_train_image_ids = {img['id'] for img in new_train_images}
            new_val_image_ids = {img['id'] for img in new_val_images}
            
            for annotation in self.all_annotations:
                if annotation['image_id'] in new_train_image_ids:
                    new_train_annotations.append(annotation)
                elif annotation['image_id'] in new_val_image_ids:
                    new_val_annotations.append(annotation)
            
            # 合并训练数据（已有训练 + 新增训练）
            final_train_images = existing_train_images + new_train_images
            final_train_annotations = existing_train_annotations + new_train_annotations
            
            # 合并验证数据（已有验证 + 新增验证）
            final_val_images = existing_val_images + new_val_images
            final_val_annotations = existing_val_annotations + new_val_annotations
            
        else:
            # 没有新数据，只使用已有数据
            final_train_images = existing_train_images
            final_train_annotations = existing_train_annotations
            final_val_images = existing_val_images
            final_val_annotations = existing_val_annotations
            print("✓ 没有新数据，保持原有数据不变")
        
        # 创建annotations目录
        annotations_dir = os.path.join(output_path, "annotations")
        os.makedirs(annotations_dir, exist_ok=True)
        
        # 生成训练数据文件
        train_data = {
            "info": self.coco_data["info"].copy(),
            "licenses": self.coco_data["licenses"].copy(),
            "images": final_train_images,
            "annotations": final_train_annotations,
            "categories": self.coco_data["categories"].copy(),
            "type": self.coco_data["type"]
        }
        
        train_file = os.path.join(annotations_dir, "instance_train.json")
        with open(train_file, 'w', encoding='utf-8') as f:
            json.dump(train_data, f, ensure_ascii=False, indent=2)
        
        # 生成验证数据文件
        val_data = {
            "info": self.coco_data["info"].copy(),
            "licenses": self.coco_data["licenses"].copy(),
            "images": final_val_images,
            "annotations": final_val_annotations,
            "categories": self.coco_data["categories"].copy(),
            "type": self.coco_data["type"]
        }
        
        val_file = os.path.join(annotations_dir, "instance_val.json")
        with open(val_file, 'w', encoding='utf-8') as f:
            json.dump(val_data, f, ensure_ascii=False, indent=2)
        
        print(f"✓ 训练数据已保存到: {train_file}")
        print(f"  - 图片数量: {len(final_train_images)} (已有: {len(existing_train_images)}, 新增: {len(final_train_images) - len(existing_train_images)})")
        print(f"  - 标注数量: {len(final_train_annotations)} (已有: {len(existing_train_annotations)}, 新增: {len(final_train_annotations) - len(existing_train_annotations)})")
        
        print(f"✓ 验证数据已保存到: {val_file}")
        print(f"  - 图片数量: {len(final_val_images)} (已有: {len(existing_val_images)}, 新增: {len(final_val_images) - len(existing_val_images)})")
        print(f"  - 标注数量: {len(final_val_annotations)} (已有: {len(existing_val_annotations)}, 新增: {len(final_val_annotations) - len(existing_val_annotations)})")
    
    def process_dataset(self, input_path: str, output_path: str):
        """处理完整的数据集生成流程"""
        print("=" * 60)
        print("开始自动化数据集生成流程")
        print("=" * 60)
        
        # 1. 加载现有数据（如果存在）
        self._load_existing_data(output_path)
        
        # 2. 检查是否为特定路径结构（handwriting/printer）
        if self._check_specific_path_structure(input_path):
            print("\n检测到特定目录结构，使用特定路径标注逻辑")
            self._process_specific_path_annotation(input_path, output_path)
            print("\n" + "=" * 60)
            print("数据集生成完成！")
            print(f"输出目录: {output_path}")
            print("=" * 60)
            return
        
        # 3. 使用PaddleX版面检测流程
        print("\n使用PaddleX版面检测流程")
        
        # 4. 初始化PaddleX版面检测模型
        if not self.initialize_layout_model():
            print("错误: PaddleX版面检测模型初始化失败，无法继续")
            return
        
        # 5. 创建输出目录结构
        images_dir = os.path.join(output_path, "images")
        os.makedirs(images_dir, exist_ok=True)
        
        # 6. 拷贝图片文件
        print("\n步骤1: 拷贝图片文件")
        image_files = self.copy_image_files(input_path, images_dir)
        
        # 7. 处理文档文件
        print("\n步骤2: 处理文档文件")
        doc_files = self.process_documents(input_path, images_dir)
        
        # 8. 合并所有图片文件
        all_image_files = image_files + doc_files
        print(f"\n总计图片文件: {len(all_image_files)}")
        
        if not all_image_files:
            print("错误: 没有找到任何图片文件")
            return
        
        # 9. 使用PaddleX进行版面检测和标注
        print("\n步骤3: 使用PaddleX进行版面检测")
        self.process_images_with_ocr(all_image_files, output_path)
        
        # 10. 生成训练和验证数据分割
        print("\n步骤4: 生成训练和验证数据分割")
        self.generate_train_val_split(output_path)
        
        print("\n" + "=" * 60)
        print("数据集生成完成！")
        print(f"输出目录: {output_path}")
        print("=" * 60)


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='自动化数据集生成脚本')
    parser.add_argument('--input_path', 
                       required=True,
                       help='源文件路径，包含需要处理的文件')
    parser.add_argument('--output_path', 
                       required=True,
                       help='输出数据集路径')
    parser.add_argument('--train_ratio', 
                       type=float,
                       default=0.8,
                       help='训练数据比例 (默认: 0.8)')
    parser.add_argument('--model_name', 
                       default='PP-DocLayout-L',
                       choices=['PP-DocLayout_plus-L', 'PP-DocLayout-L'],
                       help='版面检测模型名称 (默认: PP-DocLayout-L)')
    parser.add_argument('--model_dir', 
                       default='seal_env/PaddleX/output/best_model/inference',
                       help='模型权重地址 (默认: output/best_model/inference)')
    
    args = parser.parse_args()
    
    # 检查输入路径是否存在
    if not os.path.exists(args.input_path):
        print(f"错误: 输入路径不存在: {args.input_path}")
        return
    
    # 检查必要的库是否安装
    if not PADDLEX_AVAILABLE:
        print("错误: PaddleX未安装，请运行: pip install paddlex")
        return
    
    if not PIL_AVAILABLE:
        print("错误: PIL未安装，请运行: pip install Pillow")
        return
    
    # 创建生成器并处理数据
    generator = AutoDatasetGenerator(train_ratio=args.train_ratio, model_name=args.model_name, model_dir=args.model_dir)
    generator.process_dataset(args.input_path, args.output_path)


if __name__ == "__main__":
    main()
